import numpy as np

def create_synthetic_images():
    # 第一帧：5x5中心区域亮
    frame1 = np.array([
        [1, 1, 1, 1, 1],
        [1, 2, 2, 2, 1],
        [1, 3, 3, 3, 1],
        [1, 9, 9, 9, 1],
        [1, 5, 5, 5, 1]
    ], dtype=np.float32)
    
    # 第二帧：中心亮区向右上移动1像素
    frame2 = np.array([
        [0, 0, 0, 0, 0],
        [0, 1, 1, 1, 1],
        [0, 1, 2, 2, 2],
        [0, 1, 3, 3, 3],
        [0, 1, 9, 9, 9]
    ], dtype=np.float32)
    
    return frame1, frame2


# def lucas_kanade(image1, image2, x, y, window_size=3):
#     # 
#     Iy, Ix = np.gradient(image2)
#     p = int(window_size/2)-1
#     q = int(window_size/2)-1
    
#     It = image2 - image1
    
#     a = x-p; b = x+p+1; c=y-q; d=y+q+1
#     print(f"x={x}, y={y}, p={p}, q={q}")
#     print(range(a, b), range(c, d))
#     print(f"Ix={Ix}, \nIy={Iy}, \nIt={It}")

#     # A = Ix * Ix
#     # B = Ix * Iy
#     # C = Iy * Iy
#     # sumA = np.sum(A[a:b, c:d])
#     # sumB = np.sum(B[a:b, c:d])
#     # sumC = np.sum(C[a:b, c:d])
#     # sumD = np.sum(Ix[a:b, c:d])
#     # sumE = np.sum(Iy[a:b, c:d])
#     # sumF = np.sum(image1[a:b, c:d]-image2[a:b, c:d])
    
    
#     print(f"Ix[{a}:{b}, {c}:{d}]={Ix[a:b, c:d]}, Iy[{a}:{b}, {c}:{d}]={Iy[a:b, c:d]}, It[{a}:{b}, {c}:{d}]={It[a:b, c:d]}")
#     sumA = np.sum(Ix[a:b, c:d] ** 2)       # Ix²的和
#     sumB = np.sum(Ix[a:b, c:d] * Iy[a:b, c:d])  # IxIy的和
#     sumC = np.sum(Iy[a:b, c:d] ** 2)       # Iy²的和
#     sumD = np.sum(Ix[a:b, c:d] * It[a:b, c:d])  # Ix·It的和
#     sumE = np.sum(Iy[a:b, c:d] * It[a:b, c:d])  # Iy·It的和
    
#     print(f"sumA={sumA}, sumB={sumB}, sumC={sumC}, sumD={sumD}, sumE={sumE}")

#     denominator = sumA * sumC - sumB**2
#     if denominator == 0:
#         print(f"denominator is {denominator}")
#         return np.array([0, 0])
    
    
#     U = (sumB*sumE - sumC*sumD) / (sumA*sumC - sumB*sumB)
#     V = (sumA*sumE - sumB*sumD) / (sumA*sumC - sumB*sumB)
    
#     flow = np.array([U, V])

#     return flow


def lucas_kanade(image1, image2, x, y, window_size=1):
    # 计算空间梯度（Ix, Iy）[[9]]
    Iy, Ix = np.gradient(image2)
    # 计算时间梯度（It）[[3]]
    It = image2 - image1
    
    
    # 计算窗口半径
    p = window_size // 2
    q = window_size // 2
    
    # 提取邻域窗口内的梯度值（[[7]]）
    Ix_window = Ix[x-p:x+p+1, y-q:y+q+1]
    Iy_window = Iy[x-p:x+p+1, y-q:y+q+1]
    It_window = It[x-p:x+p+1, y-q:y+q+1]
    
    a = x-p; b = x+p+1; c=y-q; d=y+q+1
    print(f"Ix[{a}:{b}, {c}:{d}]={Ix[a:b, c:d]}, Iy[{a}:{b}, {c}:{d}]={Iy[a:b, c:d]}, It[{a}:{b}, {c}:{d}]={It[a:b, c:d]}")
    
    # 构建光流方程组（[[1]]）
    A = Ix_window.flatten()
    B = Iy_window.flatten()
    C = -It_window.flatten()  # 注意符号修正
    
    # 构造设计矩阵和观测向量（[[4]]）
    M = np.vstack((A, B)).T
    b = C
    
    # 最小二乘法求解（[[5]]）
    try:
        flow = np.linalg.lstsq(M, b, rcond=None)[0]
        u, v = flow[0], flow[1]
    except np.linalg.LinAlgError:
        u, v = 0, 0  # 奇异矩阵时返回零流
    
    return np.array([u, v])


if __name__ == '__main__':
    frame1, frame2 = create_synthetic_images()
    x, y = 3, 3
    flow = lucas_kanade(frame1, frame2, x, y, window_size=1)
    print(flow)
